#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
广告收入数据分析
本案例使用多元回归分析方法，分析电视媒体、广播媒体、报纸媒体的广告投入与产品销售额之间的关系。
"""

# 1. 业务理解
# 本问题所涉及的业务为分析电视媒体、广播媒体、报纸媒体的广告投入与产品销售额之间的关系。
# 该业务的主要内容是建立电视媒体、广播媒体、报纸媒体的广告投入与产品销售额的多元回归模型。

# 2. 数据读入
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
import os
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import statsmodels.api as sm

# 读取数据文件
data = pd.read_csv("Advertising.csv", header=0)
print("数据集的前5行：")
print(data.head())

# 3. 数据理解
# 对数据进行探索性分析
print("\n数据集的基本信息：")
print(data.info())

print("\n数据集的统计描述：")
print(data.describe())

# 绘制TV、radio、newspaper这3个变量与sales变量之间的关系图
plt.figure(figsize=(15, 5))

# 使用Seaborn的pairplot方法绘制散点图矩阵
print("\n绘制变量之间的散点图矩阵...")
sns.pairplot(data, 
             x_vars=['TV', 'radio', 'newspaper'],
             y_vars=['sales'],
             height=7,
             aspect=0.8,
             kind='reg')
plt.savefig('pairplot.png')
plt.close()

# 绘制相关系数热力图
print("\n绘制相关系数热力图...")
plt.figure(figsize=(10, 8))
correlation = data.corr()
sns.heatmap(correlation, annot=True, cmap='coolwarm', fmt='.2f')
plt.title('相关系数热力图')
plt.savefig('correlation_heatmap.png')
plt.close()

# 4. 数据准备
# 进行多元回归分析前，准备好模型所需的特征矩阵(X)和目标向量(y)
# 本例使用drop()方法删除data数据集中的Number以及sales两列返回一个DataFrame对象Data
# 并显示Data数据集的前5行数据
print("\n准备特征矩阵和目标向量...")
X = data.drop(['Number', 'sales'], axis=1)
y = data['sales']

print("\n特征矩阵X的前5行：")
print(X.head())

print("\n目标向量y的前5个值：")
print(y.head())

# 将sales的数据类型转换为NumPy中的ndarray
sales = np.ravel(y)
print("\nsales的数据类型：", type(sales))

# 5. 模型构建
# 5.1 使用统计学方法构建模型
print("\n使用statsmodels构建多元线性回归模型...")
# 添加常数列
X_add_const = sm.add_constant(X)
# 使用OLS方法构建myModel模型
myModel = sm.OLS(y, X_add_const)
# 拟合模型
results = myModel.fit()
# 打印模型摘要
print("\n模型摘要：")
print(results.summary())

# 5.2 使用机器学习方法构建模型
print("\n使用scikit-learn构建多元线性回归模型...")
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, random_state=1, test_size=0.25)

# 查看训练数据与测试数据的数量
print(f"\n训练集数据量：{X_train.shape}")
print(f"测试集数据量：{X_test.shape}")

# 使用LinearRegression进行线性回归
linreg = LinearRegression()
model = linreg.fit(X_train, y_train)
print("\n模型训练结果：")
print(model)

# 查看回归系数
print("\n模型回归系数：")
print(model.coef_)

# 查看截距
print("\n模型截距：")
print(model.intercept_)

# 计算模型在测试集上的R-squared值
score = model.score(X_test, y_test)
print(f"\n模型在测试集上的R-squared值：{score}")

# 6. 模型预测
# 使用训练好的模型进行预测
y_pred = model.predict(X_test)
print("\n预测结果：")
print(y_pred)

# 7. 模型评价
# 计算预测结果的数量
print(f"\n预测结果的数量：{len(y_pred)}")

# 绘制预测值与实际值的对比图
plt.figure(figsize=(10, 6))
plt.plot(range(len(y_pred)), y_pred, 'b', label='predict')
plt.plot(range(len(y_pred)), y_test, 'r', label='test')
plt.legend(loc='upper right')
plt.xlabel('样本序号')
plt.ylabel('销售额')
plt.title('预测值与实际值对比图')
plt.savefig('prediction_comparison.png')
plt.close()

# 计算均方误差(MSE)
from sklearn.metrics import mean_squared_error
mse = mean_squared_error(y_test, y_pred)
print(f"\n均方误差(MSE)：{mse}")

# 计算平均绝对误差(MAE)
from sklearn.metrics import mean_absolute_error
mae = mean_absolute_error(y_test, y_pred)
print(f"\n平均绝对误差(MAE)：{mae}")

# 计算R方值
from sklearn.metrics import r2_score
r2 = r2_score(y_test, y_pred)
print(f"\nR方值：{r2}")

# 结论分析
print("\n结论分析：")
print("1. 从相关系数热力图可以看出，TV广告投入与销售额的相关性最强。")
print("2. 从回归系数可以看出各个媒体广告投入对销售额的影响程度。")
print("3. 模型的R方值较高，说明该模型拟合效果较好。")
print("4. 从预测值与实际值的对比图可以直观地看出模型的预测效果。")

print("\n分析完成！可以查看生成的图表文件：pairplot.png, correlation_heatmap.png, prediction_comparison.png")
